import pytesseract
from PIL import Image
import cv2
import os
from openai import OpenAI
import re


class OCRToOpenAIProcessor:
    def __init__(self, api_key, base_url):
        self.client = OpenAI(api_key=api_key, base_url=base_url)

    def extract_text(self, image_path):
        """改进的OCR提取文本"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        # 图像尺寸调整 - 确保足够大的分辨率
        height, width = img.shape[:2]
        if max(height, width) < 1000:
            scale = 1500 / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)

        # 多种预处理方法尝试
        processed_images = []

        # 方法1: 标准灰度 + Otsu二值化
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, thresh1 = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('otsu', thresh1))

        # 方法2: 自适应阈值
        thresh2 = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                        cv2.THRESH_BINARY, 11, 2)
        processed_images.append(('adaptive', thresh2))

        # 方法3: 对比度增强
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)
        _, thresh3 = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('clahe', thresh3))

        # 尝试多种预处理结果，选择最好的
        best_text = ""
        best_confidence = 0

        for method_name, processed_img in processed_images:
            pil_img = Image.fromarray(processed_img)
            # 使用更详细的配置获取置信度信息
            text = pytesseract.image_to_string(pil_img, config=r'--oem 1 --psm 6 -l chi_sim+eng')

            # 简单评估文本质量（基于字符长度和中文字符比例）
            chinese_chars = len(re.findall(r'[\u4e00-\u9fff]', text))
            total_chars = len(text.strip())
            confidence = chinese_chars * 2 + total_chars if total_chars > 0 else 0

            if confidence > best_confidence:
                best_confidence = confidence
                best_text = text.strip()

        return best_text

    def clean_with_openai(self, ocr_text):
        """使用OpenAI清理和格式化"""
        prompt = f"""
            请分析以下从图片中识别出的题目内容。图片识别可能存在错误，请修正识别错误并提取以下信息：
            
            1. 题目描述（理论背景等）
            2. 题目本身
            3. 所有选项（A、B、C、D等）
            4. 正确答案
            
            要求：
            - 修正OCR识别中的错误字符
            
            识别内容:
            {ocr_text}
            
            请按以下格式输出，不关心分析修改的内容直接按我要求的输出：
             
            【题目描述】
            ...
            【题目】...
            【答案】...
            【选项】
            A. ...
            B. ...
            C. ...
            D. ...
            
            """

        response = self.client.chat.completions.create(
            model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
            messages=[{'role': 'user', 'content': prompt}],
            stream=True
        )

        print("OpenAI处理中...")
        full_response = ""

        for chunk in response:
            if not chunk.choices:
                continue
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                print(content, end="", flush=True)
                full_response += content
            if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                reasoning = chunk.choices[0].delta.reasoning_content
                print(reasoning, end="", flush=True)
                full_response += reasoning

        return full_response

    def extract_final_content(self, openai_response):
        """提取'下面是输出结果:'之后的内容"""
        separator = "【题目描述】"

        if separator in openai_response:
            # 找到分隔符位置
            separator_index = openai_response.find(separator)
            # 提取分隔符之后的内容
            final_content = openai_response[separator_index + len(separator):].strip()
            return "【题目描述】"+final_content
        else:
            # 如果没有找到分隔符，返回原始内容
            print("警告：未找到'下面是输出结果:'分隔符，返回完整内容")
            return openai_response

    def save_final_content(self, final_content, image_path, output_dir="final_results"):
        """保存最终内容到文件"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        base_name = os.path.splitext(os.path.basename(image_path))[0]
        output_file = os.path.join(output_dir, f"{base_name}_final.txt")

        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(final_content)

        print(f"\n✅ 最终内容已保存到: {output_file}")
        return output_file

    def process_single_image(self, image_path):
        """处理单张图片"""
        print(f"📷 处理图片: {os.path.basename(image_path)}")

        # 1. OCR识别
        print("🔍 OCR识别中...")
        ocr_text = self.extract_text(image_path)
        if not ocr_text:
            print("❌ OCR识别失败")
            return None


        print(f"📝 OCR原始内容: {ocr_text}")

        # 2. OpenAI处理
        openai_response = self.clean_with_openai(ocr_text)

        # 3. 提取最终内容
        print("\n🎯 提取最终内容...")
        final_content = self.extract_final_content(openai_response)

        # 4. 保存最终内容
        saved_file = self.save_final_content(final_content, image_path)

        return {
            'original_ocr': ocr_text,
            'openai_response': openai_response,
            'final_content': final_content,
            'saved_file': saved_file
        }


def main():
    # 配置API
    API_KEY = "sk-msxhrypvitaetkpthwdwvmmoekiabedcjxlbsznqzmjigpnq"  # 替换为您的API密钥
    BASE_URL = "https://api.siliconflow.cn/v1"  # 替换为您的API地址

    processor = OCRToOpenAIProcessor(api_key=API_KEY, base_url=BASE_URL)

    print("请选择处理模式:")
    print("1. 处理单张图片")
    print("2. 批量处理图片文件夹")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        image_path = input("请输入图片路径: ").strip()
        if not os.path.exists(image_path):
            print("❌ 文件不存在!")
            return

        result = processor.process_single_image(image_path)
        if result:
            print(f"\n🎉 处理完成!")
            print(f"📁 最终内容保存位置: {result['saved_file']}")

    elif choice == "2":
        folder_path = input("请输入图片文件夹路径: ").strip()
        if not os.path.exists(folder_path):
            print("❌ 文件夹不存在!")
            return

        # 获取所有图片文件
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']
        image_files = []
        for file in os.listdir(folder_path):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(folder_path, file))

        print(f"📂 找到 {len(image_files)} 张图片")
        
        # 询问是否自动执行
        auto_execute = input("是否自动执行所有图片处理？(y/n): ").strip().lower()
        auto_mode = auto_execute == 'y' or auto_execute == 'yes'

        results = []
        for i, image_path in enumerate(image_files, 1):
            print(f"\n{'=' * 60}")
            print(f"🔄 处理第 {i}/{len(image_files)} 张: {os.path.basename(image_path)}")
            print('=' * 60)

            result = processor.process_single_image(image_path)
            if result:
                results.append(result)

            # 如果不是最后一张且不是自动模式，等待用户确认继续
            if i < len(image_files) and not auto_mode:
                input("\n⏳ 按Enter键继续处理下一张图片...")
            elif i < len(image_files) and auto_mode:
                print("\n⏳ 自动处理下一张图片...")

        print(f"\n🎊 批量处理完成! 成功处理 {len(results)}/{len(image_files)} 张图片")

        # 生成汇总报告
        if results:
            summary_file = os.path.join("final_results", "processing_summary.txt")
            with open(summary_file, 'w', encoding='utf-8') as f:
                f.write("图片处理汇总报告\n")
                f.write("=" * 50 + "\n\n")
                f.write(f"总图片数: {len(image_files)}\n")
                f.write(f"成功处理: {len(results)}\n\n")
                f.write("处理结果文件:\n")
                for result in results:
                    f.write(f"- {os.path.basename(result['saved_file'])}\n")

            print(f"📋 汇总报告: {summary_file}")

    else:
        print("❌ 无效选择!")


# 简化版本 - 直接处理单张图片
def quick_process(image_path, api_key, base_url):
    """快速处理单张图片"""
    processor = OCRToOpenAIProcessor(api_key=api_key, base_url=base_url)

    print(f"🚀 快速处理: {os.path.basename(image_path)}")

    # OCR识别
    ocr_text = processor.extract_text(image_path)
    if not ocr_text:
        print("❌ OCR识别失败")
        return None

    # OpenAI处理
    openai_response = processor.clean_with_openai(ocr_text)

    # 提取最终内容
    final_content = processor.extract_final_content(openai_response)

    # 保存最终内容
    saved_file = processor.save_final_content(final_content, image_path)

    print(f"\n✅ 处理完成! 最终内容保存到: {saved_file}")

    return final_content


if __name__ == "__main__":
    main()